#include <stdlib.h>
#include "opencv2/highgui/highgui.hpp"
#include "opencv2/imgproc/imgproc.hpp"
#include <opencv2/highgui/highgui_c.h>
#include <thread>
#include <mutex>
#include <unistd.h>
#include <matio.h>
#include <buola/utility/random.h>
#include <buola/app/ccmdline.h>
#include <sys/stat.h>
#include <buola/tty/input.h>
#include <buola/datetime/clocks.h>

std::mutex sMutex;

using namespace cv;
using namespace buola;

struct CFrame
{
    int mIndex;
    int64_t mTimeStamp;
    Mat mDisparity;
    Mat mImage;
};

void save_frame(const std::string &pFolder,CFrame &pFrame)
{
    mkdir(pFolder.c_str(),0755);

    // save to file
    std::string lFileName=pFolder+"/image_"+to_string(pFrame.mIndex,5,'0')+".png";
    imwrite(lFileName,pFrame.mImage);
    lFileName=pFolder+"/disp_"+lNumber+".dat";
    FILE *fout = fopen(lFileName.c_str(),"wb");
    fwrite(pFrame.mDisparity.data, sizeof(float), 640*480, fout);
    fclose(fout);

    lFileName=pFolder+"/info_"+lNumber+".txt";
    io::CTextWriter lInfoFile(io::file(lFileName));
    lInfoFile << pFrame.mIndex << "\n" << pFrame.mTimeStamp << "\n";
}

int main( int argc, char* argv[] )
{
    buola_init(argc,argv);
    
    if(cmd_line().ArgCount()!=1)
    {
        msg_info() << "usage: grabsensors folder\n";
        return 1;
    }
    
    VideoCapture capture;
    capture.open(CV_CAP_OPENNI);

    if (!capture.isOpened())
        msg_info() << "capture open failed\n";
    
    capture.set(CV_CAP_PROP_OPENNI_REGISTRATION,1);

    int lFrameIndex=1;

    for(int lIter=0;lIter<50;lIter++)
    {
        msg_info() << "\n\n\npress key to start\n";
        tty::get_char();
        msg_info() << "\n\n\n";
            
        lFrameIndex=1;

        std::vector<CFrame*> lFrames;

        while(true)
        {              
            CFrame *lFrame=new CFrame;
            Mat lDepthMap,lShowDepth,lImage,lDisparity;
                    
            capture.grab();
            capture.retrieve( lDepthMap, CV_CAP_OPENNI_DEPTH_MAP );
            capture.retrieve( lImage, CV_CAP_OPENNI_BGR_IMAGE );
            capture.retrieve( lDisparity, CV_CAP_OPENNI_DISPARITY_MAP_32F );
            lFrame->mImage=lImage.clone();
            lFrame->mDisparity=lDisparity.clone();

            lDepthMap.convertTo(lShowDepth, CV_8UC1, 255.0/5000.0);
                    
//            imshow("bgr",lImage);
//            imshow("depth",lShowDepth);
            lFrame->mIndex=lFrameIndex++;
            lFrame->mTimeStamp=CCalendarClock::now().time_since_epoch().count();
            lFrames.push_back(lFrame);
            
            if(tty::can_get_char()) break;
//            waitKey(1);
        }
        
        tty::get_char();
        
        mkdir(cmd_line().Arg(0).c_str(),0755);
        std::string lFolder=cmd_line().Arg(0)/("set_"+to_string(lIter));

        for(CFrame *lFrame : lFrames)
        {
            save_frame(lFolder,*lFrame);
            delete lFrame;
        }
    }

    return buola_finish();
}

